"""
使用pytorch进行训练
f(x)=x*w
loss=(f(xn)-y)**2
"""
import random
import matplotlib.pyplot as plt
import torch

# 定义已知的数据集,training的dataset
x_data = [1.0, 2.0, 3.0]
y_data = [2.0, 4.0, 6.0]

"""
关于torch的tensor张量

这些属性得到的都是张量tensor:
data属性储存权重w
grad属性储存梯度(loss对w的倒数)

这些属性得到的都是标量:
item获得当前的标量值
"""

# 定义权重w(张量)
w = torch.Tensor([1.0])
# 声明需要计算梯度
w.requires_grad = True

# 定义前馈
def forward(x):
    """
    前馈
    :param x: 输入的x值,这里因为w是tensor张量,所以会自动转换为张量
    :return: 前馈张量,是一个tensor
    """
    return x * w


# 单个样本的损失
def loss(x, y):
    # 计算y_hat的计算图
    y_hat = forward(x)
    # 计算单个样本的损失计算图
    return (y_hat - y) ** 2


# 记录损失
allLoss = []
# 记录轮数
allEpoch = []

# 进行100轮训练
for epoch in range(100):
    # 记录
    allEpoch.append(epoch)
    # 使用随机单样本
    idx = random.randint(0, len(x_data) - 1)
    x = x_data[idx]
    y = y_data[idx]
    # 构建单样本损失图(前馈计算loss)
    tempLoss = loss(x, y)
    # 记录
    allLoss.append(tempLoss.item())
    # 计算反向传播,算出梯度
    tempLoss.backward()
    # 权重更新,学习率0.01*梯度
    w.data -= 0.01 * w.grad.data
    # 将梯度的计算图置零,否则在循环中会一直进行梯度叠加
    w.grad.data.zero_()
    print("progress:", epoch, "w: ", w.item(), "loss:", tempLoss.item())

print("the value of 4 is:", forward(4).item())

# 绘制损失-轮数图
plt.plot(allEpoch, allLoss)
plt.ylabel("loss")
plt.xlabel("epoch")
plt.grid()
plt.show()
